package com.gxc.tree;

/**
 * 前缀树（字典树）
 */
public class Trie {

    private TrieNode root;

    public Trie() {
        this.root = new TrieNode();
    }

    public void insertStr(String word) {
        if (word == null) return;

        char[] chars = word.toCharArray();

        TrieNode trieNode = this.root;
        //记录总共有多少word
        trieNode.pass++;
        int pos = 0;
        for (int i = 0; i < chars.length; i++) {
            char c = chars[i];
            pos = c - 'a';

            if (trieNode.next[pos] == null) {
                trieNode.next[pos] = new TrieNode();
            }
            trieNode = trieNode.next[pos];
            trieNode.pass++;
        }
        trieNode.end++;
    }

    //word加入过几次
    public int search(String word) {
        if (word == null) return 0;

        char[] chars = word.toCharArray();

        TrieNode trieNode = this.root;
        int pos = 0;
        for (char c : chars) {
            pos = c - 'a';
            if (trieNode.next[pos] == null) return 0;

            trieNode = trieNode.next[pos];
        }
        return trieNode.end;
    }

    //以word为前缀的总共有多少
    public int prefixNumber(String word) {
        if (word == null) return 0;

        char[] chars = word.toCharArray();

        TrieNode trieNode = this.root;
        int pos = 0;
        for (char c : chars) {
            pos = c - 'a';
            if (trieNode.next[pos] == null) return 0;

            trieNode = trieNode.next[pos];
        }
        return trieNode.pass;
    }

    public void deleteStr(String word) {
        if (word == null) return;

        if (search(word)>0) {
            char[] chars = word.toCharArray();

            TrieNode trieNode = this.root;
            trieNode.pass--;

            int pos = 0;
            for (char c : chars) {
                pos = c - 'a';
                if (--trieNode.next[pos].pass == 0) {
                    trieNode.next[pos] = null;
                    return;
                }
                trieNode = trieNode.next[pos];
            }
            trieNode.end--;
        }
    }
}
